import os.path
import shutil
import sys
import logging
import requests
import time
from datetime import datetime, timedelta

import akshare as ak
import numpy as np
import pandas as pd

# 配置日志格式，包含文件名和行号
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(filename)s:%(lineno)d - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# stock_sse_summary_df = ak.stock_sse_summary()
# print(stock_sse_summary_df)


# stock_sse_deal_daily_df = ak.stock_sse_deal_daily(date="20211227")
# print(stock_sse_deal_daily_df)


# stock_zh_a_spot_em_df = ak.stock_zh_a_spot_em()
# print(stock_zh_a_spot_em_df)
# stock_zh_a_spot_em_df.to_excel("行情沪深京A.xlsx", sheet_name="沪深京A",index=False)


# stock_industry_category_cninfo_df = ak.stock_industry_category_cninfo(symbol="巨潮行业分类标准")
# print(stock_industry_category_cninfo_df)
# stock_industry_category_cninfo_df.to_excel("行业分类.xlsx",sheet_name="巨潮",index=False)

# stock_info_a_code_name_df = ak.stock_info_a_code_name()
# print(stock_info_a_code_name_df)
# stock_info_a_code_name_df.to_excel("")

# 行业市盈率
# stock_industry_pe_ratio_cninfo_df = ak.stock_industry_pe_ratio_cninfo(symbol="国证行业分类", date="20240617")
# print(stock_industry_pe_ratio_cninfo_df)
# stock_industry_pe_ratio_cninfo_df.to_excel("行业市盈率.xlsx",sheet_name="行业市盈率",index=False)

code = "601633"
name = "长城汽车"

code = "002594"
name = "比亚迪"

code = "300750"
name = "宁德时代"

code = "601808"
name = "中海油服"

# stock_a_indicator_lg_df = ak.stock_a_indicator_lg(symbol=code)
# stock_a_indicator_lg_df.to_excel(name+".xlsx", sheet_name=name, index=False)


# 000016	上证50

# stock_zh_valuation_baidu_df = ak.stock_zh_valuation_baidu(symbol=code, indicator="市盈率(TTM)", period="近一年")
# stock_zh_valuation_baidu_df.to_excel(name+".xlsx", sheet_name=name, index=False)

# 股票指数信息一览表
# index_stock_info_df = ak.index_stock_info()
# print(index_stock_info_df)

# stock_profit_forecast_em_df = ak.stock_profit_forecast_em()
# print(stock_profit_forecast_em_df)

code = "000905"
name = "中证500"


def current_day():
    now = datetime.now()
    date_str = now.strftime('%Y%m%d')
    return date_str


def years_ago(year=5):
    now = datetime.now()
    five_years_ago = now - timedelta(days=year * 365)
    date_str = five_years_ago.strftime('%Y%m%d')
    return date_str


def percentile(data):
    idx = 0
    f = data.iloc[0]
    for d in data:
        if 0 < d < f:
            idx = idx + 1
    return round(idx * 1.0 / len(data), 3)


def day_count(col_day, year=1):
    year_ago_str = years_ago(year=year)
    cnt = 0
    for day in col_day:
        day_str = day.strftime("%Y%m%d")
        if day_str >= year_ago_str:
            cnt = cnt + 1
    return cnt


def expand_list_with_list_comprehension(lst, size, fill_value=None):
    if len(lst) >= size:
        return lst[:size]
    return lst + [fill_value] * (size - len(lst))


def cal_percentile(col_day, col_data, year=1):
    year1 = day_count(col_day, year=year)
    pe_ttm_val = []
    for i in range(year1):
        pe_ttm = col_data[i:i + year1]
        val = percentile(pe_ttm)
        pe_ttm_val.append(val)
    return expand_list_with_list_comprehension(pe_ttm_val, len(col_data))


def is_low(pe_ttm, delta=0.05):
    if pe_ttm.iloc[0] < delta:
        return True


def str_ttm_year123(pe_ttm1, pe_ttm2, pe_ttm3):
    if pe_ttm1 is not None and pe_ttm2 is not None and pe_ttm3 is not None:
        return "%.2f-%.2f-%.2f" % (pe_ttm1.iloc[0], pe_ttm2.iloc[0], pe_ttm3.iloc[0])
    return ""


def ext_name_with_pettm(pe_ttm1, pe_ttm2, pe_ttm3, name):
    ttm_year123 = str_ttm_year123(pe_ttm1, pe_ttm2, pe_ttm3)
    if ttm_year123:
        ext_name = "indicator-%s-%s.xlsx" % (ttm_year123, name)
    else:
        ext_name = "indicator-%s.xlsx" % (name)
    return ext_name


# 600115 - 中国东航

def is_nan(pe_ttm, size=5):
    lst = [d for d in pe_ttm.iloc[0:size]]
    # logger.debug(f"数据列表: {lst}, 是否有NaN: {np.any(np.isnan(lst))}")
    return np.isnan(lst[0]) or np.any(np.isnan(lst)) or all(not x for x in lst)

def to_txt(path, max, min, diff):
    lines = [
        "max,%s\n" % max,
        "min,%s\n" % min,
        "max/min,%.2f\n" % diff
    ]
    with open(path+".txt", "w") as f:
        f.writelines(lines)

def stock_indicator_evaluate(out_dir, code="000300", name="沪深300"):
    logger.info(f"开始评估股票指标: {code} - {name}")
    
    try:
        # 获取股票指标数据 - 添加重试机制
        max_retries = 3
        retry_count = 0
        stock_a_indicator_lg_df = None
        
        while retry_count < max_retries:
            try:
                stock_a_indicator_lg_df = ak.stock_a_indicator_lg(symbol=code)
                if stock_a_indicator_lg_df is not None and not stock_a_indicator_lg_df.empty:
                    break
            except (requests.exceptions.JSONDecodeError, ValueError) as json_error:
                retry_count += 1
                if retry_count >= max_retries:
                    raise
                logger.warning(f"获取股票指标数据时出现JSON解码错误，正在重试 ({retry_count}/{max_retries}): {code} - {name}")
                time.sleep(1)  # 添加延迟后重试
        
        logger.debug(f"成功获取股票 {code} 的指标数据，共 {len(stock_a_indicator_lg_df)} 条记录")
        
        # 检查必要列是否存在
        if "trade_date" not in stock_a_indicator_lg_df.columns or "pe_ttm" not in stock_a_indicator_lg_df.columns:
            logger.error(f"股票 {code} 的指标数据缺少必要列")
            return
        
        df = stock_a_indicator_lg_df.sort_values(by="trade_date", ascending=False)
        col_day = df["trade_date"]
        col_pe_ttm = df["pe_ttm"]
        all_pe_ttm_nan = is_nan(col_pe_ttm, size=10)

        if not os.path.exists(out_dir):
            logger.info(f"创建输出目录: {out_dir}")
            os.makedirs(out_dir)

        try:
            # 获取历史行情数据
            logger.debug(f"获取股票 {code} 的历史行情数据")
            stock_zh_a_hist_df = ak.stock_zh_a_hist(symbol=code, period="daily", 
                                                  start_date=years_ago(year=10),
                                                  end_date=current_day(), adjust="qfq")

            if stock_zh_a_hist_df is not None and not stock_zh_a_hist_df.empty and "日期" in stock_zh_a_hist_df.columns and "收盘" in stock_zh_a_hist_df.columns:
                df_hist = stock_zh_a_hist_df.sort_values(by="日期", ascending=False)
                col_sp = df_hist["收盘"]
                ext_name = "%s-hist.xlsx" % name
                p = os.path.join(out_dir, ext_name)
                df_hist.to_excel(p, sheet_name=name, index=False)
                logger.info(f"保存历史行情数据到: {p}")

                df["收盘"] = expand_list_with_list_comprehension(list(col_sp.values), len(col_pe_ttm))
            else:
                logger.warning(f"无法获取股票 {code} 的历史行情数据或数据不完整")
                df["收盘"] = [None] * len(df)
        except Exception as hist_error:
            logger.error(f"获取历史行情数据时出错: {code} - {name} - {str(hist_error)}")
            df["收盘"] = [None] * len(df)

        ext_name = "indicator-%s.xlsx" % name

        if not all_pe_ttm_nan:
            try:
                logger.debug(f"计算股票 {code} 的PE百分位数据")
                df["pe_ttm_year1"] = cal_percentile(col_day, col_pe_ttm, year=1)
                df["pe_ttm_year2"] = cal_percentile(col_day, col_pe_ttm, year=2)
                df["pe_ttm_year3"] = cal_percentile(col_day, col_pe_ttm, year=3)
                df["pe_ttm_year4"] = cal_percentile(col_day, col_pe_ttm, year=4)
                df["pe_ttm_year5"] = cal_percentile(col_day, col_pe_ttm, year=5)
                ext_name = ext_name_with_pettm(df["pe_ttm_year1"], df["pe_ttm_year2"], df["pe_ttm_year3"], name)
                
                # 保存文本数据（如果有收盘数据）
                try:
                    if "收盘" in df.columns and df["收盘"].notna().any():
                            收盘_values = df["收盘"].dropna().values
                            if len(收盘_values) > 0:
                                txt_path = os.path.join(out_dir, ext_name)
                                to_txt(txt_path, 收盘_values.max(), 收盘_values.min(), 收盘_values.max() / 收盘_values.min())
                                logger.debug(f"保存文本数据到: {txt_path}.txt")
                except Exception as txt_error:
                    logger.error(f"保存文本数据时出错: {code} - {name} - {str(txt_error)}")
            except Exception as calc_error:
                logger.error(f"计算PE百分位数据时出错: {code} - {name} - {str(calc_error)}")

        # 确保保存数据，即使部分计算失败
        try:
            p = os.path.join(out_dir, ext_name)
            df.to_excel(p, sheet_name=name, index=False)
            logger.info(f"保存指标评估数据到: {p}")
        except Exception as save_error:
            logger.error(f"保存数据时出错: {code} - {name} - {str(save_error)}")
            
    except requests.exceptions.JSONDecodeError as json_error:
        logger.error(f"处理股票指标时出现JSON解码错误: {code} - {name}", exc_info=True)
    except Exception as e:
        logger.error(f"处理股票指标时出错: {code} - {name} - {str(e)}", exc_info=True)

    # if all_pe_ttm_nan:
    #     return

    # for delta in [0.05, 0.08, 0.1, 0.2]:
    #     low = is_low(df["pe_ttm_year1"], delta=delta) and is_low(df["pe_ttm_year2"],
    #                                                              delta=delta) and is_low(df["pe_ttm_year3"],
    #                                                                                      delta=delta)
    #
    #     if low:
    #         # df["收盘"] = expand_list_with_list_comprehension(list(col_sp.values), len(col_pe_ttm))
    #         ext_name = "a-low-%.2f-indicator-%s.xlsx" % (delta,  name)
    #         p = os.path.join(out_dir, ext_name)
    #         df.to_excel(p, sheet_name=name, index=False)
    #         break


# stock_indicator_evaluate("out_dir", code="601360", name="三六零")
# stock_indicator_evaluate("out_dir", code="600547", name="山东黄金")


def index_stock_mmoney(code="000016", name="上证50"):
    # 股票指数成分
    logger.info(f"开始获取指数 {name}({code}) 的成分股数据")
    index_stock_cons_sina = ak.index_stock_cons_sina(symbol=code)
    out_dir = os.path.join(name, current_day())

    if os.path.exists(out_dir) and os.path.isdir(out_dir):
        logger.info(f"删除已存在的目录: {out_dir}")
        shutil.rmtree(out_dir)
    logger.info(f"创建输出目录: {out_dir}")
    os.makedirs(out_dir)

    p = os.path.join(out_dir, name + "-cons.xlsx")
    logger.info(f"保存成分股数据到: {p}")
    index_stock_cons_sina.to_excel(p, sheet_name=name, index=False)

    total_stocks = len(index_stock_cons_sina)
    logger.info(f"开始处理 {total_stocks} 只成分股")
    
    for index, row in index_stock_cons_sina.iterrows():
        stock_code, stock_name = row['code'], row['name']
        logger.info(f"处理第 {index+1}/{total_stocks} 只股票: {stock_code} - {stock_name}")

        try:
            stock_indicator_evaluate(out_dir, code=stock_code, name=stock_name)
        except Exception as e:
            logger.error(f"处理股票失败: {stock_code} - {stock_name} - {str(e)}")


def stock_summary_a(day=current_day()):
    logger.info(f"开始生成A股行情汇总，日期: {day}")
    
    try:
        # 获取上证指数数据
        logger.debug("获取上证指数历史数据")
        sh_df = ak.index_zh_a_hist(symbol="000001", period="daily", 
                                 start_date=years_ago(year=3),
                                 end_date=day)
        
        # 获取深证成指数据
        logger.debug("获取深证成指历史数据")
        sz_df = ak.index_zh_a_hist(symbol="399001", period="daily", 
                                 start_date=years_ago(year=3),
                                 end_date=day)
        
        sz_df = sz_df.drop(sz_df.columns[0], axis=1)

        # 转换成交额单位为亿元
        logger.debug("处理成交额数据，转换为亿元单位")
        sh_df["成交额"] = round(sh_df["成交额"] / (10000 * 10000), 3)
        sz_df["成交额"] = round(sz_df["成交额"] / (10000 * 10000), 3)

        # 合并数据并计算成交总额
        df = pd.concat([sh_df, sz_df], axis=1)
        df["成交总额"] = round(sh_df["成交额"] + sz_df["成交额"], 1)

        # 排序并保存
        df = df.sort_values(df.columns[0], ascending=False)
        df.to_excel("A股行情.xlsx", sheet_name="a", index=False)
        logger.info("A股行情汇总数据已保存到A股行情.xlsx")
        
    except Exception as e:
        logger.error(f"生成A股行情汇总时出错: {str(e)}", exc_info=True)


def index_indicator():
    logger.info("开始处理指数成分股指标")
    
    index_stock_codes = [
        # ("930050", "中证A50"),
        # ("000510", "中证A500"),
        ("000016", "上证50"),
        ("000300", "沪深300"),
        ("000905", "中证500"),
    ]

    for i, (code, name) in enumerate(index_stock_codes):
        logger.info(f"处理第 {i+1}/{len(index_stock_codes)} 个指数: {name}({code})")
        try:
            index_stock_mmoney(code=code, name=name)
        except Exception as e:
            logger.error(f"处理指数 {name}({code}) 时出错: {str(e)}", exc_info=True)
    
    logger.info("所有指数成分股指标处理完成")


if __name__ == "__main__":
    logger.info("程序开始执行")
    
    try:
        if len(sys.argv) >= 2 and sys.argv[1] == "a":
            logger.info("执行模式: 仅生成A股行情汇总")
            stock_summary_a()
        else:
            logger.info("执行模式: 生成A股行情汇总和指数成分股指标")
            stock_summary_a()
            index_indicator()
        
        logger.info("程序执行完成")
        
    except KeyboardInterrupt:
        logger.warning("程序被用户中断")
        sys.exit(1)
    except Exception as e:
        logger.critical(f"程序执行过程中发生严重错误: {str(e)}", exc_info=True)
        sys.exit(1)
